# encoding=utf-8
import torch
import openvino as ov
from networks import get_norm_layer, ResUNet

device = torch.device("cpu")
model = ResUNet(3, 3, 64, norm_layer=get_norm_layer())
model_parameter = torch.load("120_net_G_A.pth", map_location=device)
model.load_state_dict(model_parameter)
model.eval()

print("ov version", ov.get_version())
ov_model = ov.convert_model(model,
                            example_input=torch.rand(1, 3, 512, 512),
                            input=("x", [1, 3, 512, 512]),
                            verbose=True)
ov.save_model(ov_model, "export_funds_stillgan_cpu.xml", compress_to_fp16=False)

